{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1b280ab-b61f-4d1a-bf7e-44e5f9ed3a5c",
   "metadata": {
    "id": "e1b280ab-b61f-4d1a-bf7e-44e5f9ed3a5c"
   },
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efde77f2-6af3-4781-8597-89ecd3f41a52",
   "metadata": {
    "id": "efde77f2-6af3-4781-8597-89ecd3f41a52"
   },
   "source": [
    "# Qwen3 Mixture-of-Experts From Scratch (A Standalone Notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55cdef4d-de59-4a65-89f9-fa2a8ef3471d",
   "metadata": {
    "id": "55cdef4d-de59-4a65-89f9-fa2a8ef3471d"
   },
   "source": [
    "- This notebook is purposefully minimal and focuses on the code to implement Qwen3-30B-A3B model (with support for **Coder**, **Instruct** and **Thinking** variants); for more information about this model, please see the original blog post, technical report, and model hub pages:\n",
    "  - [Qwen3: Think Deeper, Act Faster](https://qwenlm.github.io/blog/qwen3/)\n",
    "  - [Qwen3 Technical Report](https://arxiv.org/abs/2505.09388)\n",
    "  - https://huggingface.co/Qwen/Qwen3-Coder-30B-A3B-Instruct (Qwen3 Coder Flash)\n",
    "  - https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507 (new thinking model)\n",
    "  - https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507 (new instruct model)\n",
    "  - https://huggingface.co/Qwen/Qwen3-30B-A3B (original Instruct/Thinking hybrid model)\n",
    "- Many architectural components in Qwen3 are similar to Llama 3; for a step-by-step guide that explains the individual components and the relationship between GPT and the components used here, you may like the GPT-to-Llama conversion notebooks:\n",
    "  - [Converting a From-Scratch GPT Architecture to Llama 2](../07_gpt_to_llama/converting-gpt-to-llama2.ipynb)\n",
    "  - [Converting Llama 2 to Llama 3.2 From Scratch](../07_gpt_to_llama/converting-llama2-to-llama3.ipynb)\n",
    "  \n",
    "\n",
    "**By default, this notebook runs Qwen3-Coder-30B-A3B-Instruct (aka Qwen3 Coder Flash) and requires 80 GB of VRAM (e.g., a single A100 or H100)**\n",
    "\n",
    "<br>\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/qwen/qwen3-coder-flash-overview.webp?123\" width=\"600px\">\n",
    "\n",
    "<br>\n",
    "  \n",
    "- About the code:\n",
    "  - all code is my own code, mapping the Qwen3 architecture onto the model code implemented in my [Build A Large Language Model (From Scratch)](http://mng.bz/orYv) book; the code is released under a permissive open-source Apache 2.0 license (see [LICENSE.txt](https://github.com/rasbt/LLMs-from-scratch/blob/main/LICENSE.txt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c201adb-747e-437b-9a62-442802941e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install -r https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/refs/heads/main/ch05/07_gpt_to_llama/requirements-extra.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dd1b65a8-4301-444a-bd7c-a6f2bd1df9df",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dd1b65a8-4301-444a-bd7c-a6f2bd1df9df",
    "outputId": "4f762354-e0a3-4cc2-e5d4-e61a227a202c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface_hub version: 0.34.3\n",
      "tokenizers version: 0.21.4\n",
      "torch version: 2.7.1+cu128\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\n",
    "    \"huggingface_hub\",  # to download pretrained weights\n",
    "    \"tokenizers\",       # to implement the tokenizer\n",
    "    \"torch\",            # to implement the model\n",
    "]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "653410a6-dd2b-4eb2-a722-23d9782e726d",
   "metadata": {
    "id": "653410a6-dd2b-4eb2-a722-23d9782e726d"
   },
   "source": [
    "&nbsp;\n",
    "# 1. Architecture code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "82076c21-9331-4dcd-b017-42b046cf1a60",
   "metadata": {
    "id": "82076c21-9331-4dcd-b017-42b046cf1a60"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.fc2 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.fc3 = nn.Linear(cfg[\"hidden_dim\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_fc1 = self.fc1(x)\n",
    "        x_fc2 = self.fc2(x)\n",
    "        x = nn.functional.silu(x_fc1) * x_fc2\n",
    "        return self.fc3(x)\n",
    "\n",
    "\n",
    "class MoEFeedForward(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.num_experts_per_tok = cfg[\"num_experts_per_tok\"]\n",
    "        self.num_experts = cfg[\"num_experts\"]\n",
    "        self.emb_dim = cfg[\"emb_dim\"]\n",
    "        self.gate = nn.Linear(cfg[\"emb_dim\"], cfg[\"num_experts\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "\n",
    "        self.fc1 = nn.ModuleList([nn.Linear(cfg[\"emb_dim\"], cfg[\"moe_hidden_dim\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "                                  for _ in range(cfg[\"num_experts\"])])\n",
    "        self.fc2 = nn.ModuleList([nn.Linear(cfg[\"emb_dim\"], cfg[\"moe_hidden_dim\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "                                  for _ in range(cfg[\"num_experts\"])])\n",
    "        self.fc3 = nn.ModuleList([nn.Linear(cfg[\"moe_hidden_dim\"], cfg[\"emb_dim\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "                                  for _ in range(cfg[\"num_experts\"])])\n",
    "\n",
    "    def forward(self, x):\n",
    "        scores = self.gate(x)  # (b, seq_len, num_experts)\n",
    "        topk_scores, topk_indices = torch.topk(scores, self.num_experts_per_tok, dim=-1)\n",
    "        topk_probs = torch.softmax(topk_scores, dim=-1)\n",
    "\n",
    "        batch, seq_len, _ = x.shape\n",
    "        x_flat = x.reshape(batch * seq_len, -1)\n",
    "        out_flat = torch.zeros(batch * seq_len, self.emb_dim, device=x.device, dtype=x.dtype)\n",
    "\n",
    "        topk_indices_flat = topk_indices.reshape(-1, self.num_experts_per_tok)\n",
    "        topk_probs_flat = topk_probs.reshape(-1, self.num_experts_per_tok)\n",
    "\n",
    "        unique_experts = torch.unique(topk_indices_flat)\n",
    "\n",
    "        for expert_id_tensor in unique_experts:\n",
    "            expert_id = int(expert_id_tensor.item())\n",
    "            mask = topk_indices_flat == expert_id\n",
    "            if not mask.any():\n",
    "                continue\n",
    "\n",
    "            token_mask = mask.any(dim=-1)\n",
    "            selected_idx = token_mask.nonzero(as_tuple=False).squeeze(-1)\n",
    "            if selected_idx.numel() == 0:\n",
    "                continue\n",
    "\n",
    "            expert_input = x_flat.index_select(0, selected_idx)\n",
    "            hidden = torch.nn.functional.silu(self.fc1[expert_id](expert_input)) * self.fc2[expert_id](expert_input)\n",
    "            expert_out = self.fc3[expert_id](hidden)\n",
    "\n",
    "            mask_selected = mask[selected_idx]\n",
    "            slot_indices = mask_selected.int().argmax(dim=-1, keepdim=True)\n",
    "            selected_probs = torch.gather(topk_probs_flat.index_select(0, selected_idx), dim=-1, index=slot_indices).squeeze(-1)\n",
    "\n",
    "            out_flat.index_add_(0, selected_idx, expert_out * selected_probs.unsqueeze(-1))\n",
    "\n",
    "        return out_flat.reshape(batch, seq_len, self.emb_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "56715760-37e1-433e-89da-04864c139a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, emb_dim, eps=1e-6, bias=False, qwen3_compatible=True):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.qwen3_compatible = qwen3_compatible\n",
    "        self.scale = nn.Parameter(torch.ones(emb_dim))\n",
    "        self.shift = nn.Parameter(torch.zeros(emb_dim)) if bias else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        input_dtype = x.dtype\n",
    "\n",
    "        if self.qwen3_compatible:\n",
    "            x = x.to(torch.float32)\n",
    "\n",
    "        variance = x.pow(2).mean(dim=-1, keepdim=True)\n",
    "        norm_x = x * torch.rsqrt(variance + self.eps)\n",
    "        norm_x = norm_x * self.scale\n",
    "\n",
    "        if self.shift is not None:\n",
    "            norm_x = norm_x + self.shift\n",
    "\n",
    "        return norm_x.to(input_dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4b9a346f-5826-4083-9162-abd56afc03f0",
   "metadata": {
    "id": "4b9a346f-5826-4083-9162-abd56afc03f0"
   },
   "outputs": [],
   "source": [
    "def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, dtype=torch.float32):\n",
    "    assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n",
    "\n",
    "    # Compute the inverse frequencies\n",
    "    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype)[: (head_dim // 2)].float() / head_dim))\n",
    "\n",
    "    # Generate position indices\n",
    "    positions = torch.arange(context_length, dtype=dtype)\n",
    "\n",
    "    # Compute the angles\n",
    "    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)  # Shape: (context_length, head_dim // 2)\n",
    "\n",
    "    # Expand angles to match the head_dim\n",
    "    angles = torch.cat([angles, angles], dim=1)  # Shape: (context_length, head_dim)\n",
    "\n",
    "    # Precompute sine and cosine\n",
    "    cos = torch.cos(angles)\n",
    "    sin = torch.sin(angles)\n",
    "\n",
    "    return cos, sin\n",
    "\n",
    "\n",
    "def apply_rope(x, cos, sin, offset=0):\n",
    "    # x: (batch_size, num_heads, seq_len, head_dim)\n",
    "    batch_size, num_heads, seq_len, head_dim = x.shape\n",
    "    assert head_dim % 2 == 0, \"Head dimension must be even\"\n",
    "\n",
    "    # Split x into first half and second half\n",
    "    x1 = x[..., : head_dim // 2]  # First half\n",
    "    x2 = x[..., head_dim // 2:]  # Second half\n",
    "\n",
    "    # Adjust sin and cos shapes\n",
    "    cos = cos[offset:offset + seq_len, :].unsqueeze(0).unsqueeze(0)  # Shape: (1, 1, seq_len, head_dim)\n",
    "    sin = sin[offset:offset + seq_len, :].unsqueeze(0).unsqueeze(0)\n",
    "\n",
    "    # Apply the rotary transformation\n",
    "    rotated = torch.cat((-x2, x1), dim=-1)\n",
    "    x_rotated = (x * cos) + (rotated * sin)\n",
    "\n",
    "    # It's ok to use lower-precision after applying cos and sin rotation\n",
    "    return x_rotated.to(dtype=x.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e8169ab5-f976-4222-a2e1-eb1cabf267cb",
   "metadata": {
    "id": "e8169ab5-f976-4222-a2e1-eb1cabf267cb"
   },
   "outputs": [],
   "source": [
    "class GroupedQueryAttention(nn.Module):\n",
    "    def __init__(\n",
    "        self, d_in, num_heads, num_kv_groups, head_dim=None, qk_norm=False, dtype=None\n",
    "    ):\n",
    "        super().__init__()\n",
    "        assert num_heads % num_kv_groups == 0, \"num_heads must be divisible by num_kv_groups\"\n",
    "\n",
    "        self.num_heads = num_heads\n",
    "        self.num_kv_groups = num_kv_groups\n",
    "        self.group_size = num_heads // num_kv_groups\n",
    "\n",
    "        if head_dim is None:\n",
    "            assert d_in % num_heads == 0, \"`d_in` must be divisible by `num_heads` if `head_dim` is not set\"\n",
    "            head_dim = d_in // num_heads\n",
    "\n",
    "        self.head_dim = head_dim\n",
    "        self.d_out = num_heads * head_dim\n",
    "\n",
    "        self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)\n",
    "        self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)\n",
    "        self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)\n",
    "\n",
    "        self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)\n",
    "\n",
    "        if qk_norm:\n",
    "            self.q_norm = RMSNorm(head_dim, eps=1e-6)\n",
    "            self.k_norm = RMSNorm(head_dim, eps=1e-6)\n",
    "        else:\n",
    "            self.q_norm = self.k_norm = None\n",
    "\n",
    "    def forward(self, x, mask, cos, sin, start_pos=0, cache=None):\n",
    "        b, num_tokens, _ = x.shape\n",
    "\n",
    "        # Apply projections\n",
    "        queries = self.W_query(x)  # (b, num_tokens, num_heads * head_dim)\n",
    "        keys = self.W_key(x)       # (b, num_tokens, num_kv_groups * head_dim)\n",
    "        values = self.W_value(x)   # (b, num_tokens, num_kv_groups * head_dim)\n",
    "\n",
    "        # Reshape\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        keys_new = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)\n",
    "        values_new = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # Optional normalization\n",
    "        if self.q_norm:\n",
    "            queries = self.q_norm(queries)\n",
    "        if self.k_norm:\n",
    "            keys_new = self.k_norm(keys_new)\n",
    "\n",
    "        # Apply RoPE\n",
    "        queries = apply_rope(queries, cos, sin, offset=start_pos)\n",
    "        keys_new = apply_rope(keys_new, cos, sin, offset=start_pos)\n",
    "\n",
    "        if cache is not None:\n",
    "            prev_k, prev_v = cache\n",
    "            keys = torch.cat([prev_k, keys_new], dim=2)\n",
    "            values = torch.cat([prev_v, values_new], dim=2)\n",
    "            next_cache = (keys, values)\n",
    "        else:\n",
    "            start_pos = 0  # reset RoPE\n",
    "            keys, values = keys_new, values_new\n",
    "            next_cache = (keys, values)\n",
    "\n",
    "        # Expand K and V to match number of heads\n",
    "        keys = keys.repeat_interleave(self.group_size, dim=1)\n",
    "        values = values.repeat_interleave(self.group_size, dim=1)\n",
    "\n",
    "        # Attention\n",
    "        attn_scores = queries @ keys.transpose(2, 3)\n",
    "        attn_scores = attn_scores.masked_fill(mask, -torch.inf)\n",
    "        attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)\n",
    "\n",
    "        context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)\n",
    "        return self.out_proj(context), next_cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "457cb2f8-50c1-4045-8a74-f181bfb5fea9",
   "metadata": {
    "id": "457cb2f8-50c1-4045-8a74-f181bfb5fea9"
   },
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.att = GroupedQueryAttention(\n",
    "            d_in=cfg[\"emb_dim\"],\n",
    "            num_heads=cfg[\"n_heads\"],\n",
    "            head_dim=cfg[\"head_dim\"],\n",
    "            num_kv_groups=cfg[\"n_kv_groups\"],\n",
    "            qk_norm=cfg[\"qk_norm\"],\n",
    "            dtype=cfg[\"dtype\"]\n",
    "        )\n",
    "        if cfg[\"num_experts\"] > 0:\n",
    "            self.ff = MoEFeedForward(cfg)\n",
    "        else:\n",
    "            self.ff = FeedForward(cfg)\n",
    "        self.norm1 = RMSNorm(cfg[\"emb_dim\"], eps=1e-6)\n",
    "        self.norm2 = RMSNorm(cfg[\"emb_dim\"], eps=1e-6)\n",
    "\n",
    "    def forward(self, x, mask, cos, sin, start_pos=0, cache=None):\n",
    "        # Shortcut connection for attention block\n",
    "        shortcut = x\n",
    "        x = self.norm1(x)\n",
    "        x, next_cache = self.att(x, mask, cos, sin, start_pos=start_pos, cache=cache)  # Shape [batch_size, num_tokens, emb_size]\n",
    "        x = x + shortcut  # Add the original input back\n",
    "\n",
    "        # Shortcut connection for feed-forward block\n",
    "        shortcut = x\n",
    "        x = self.norm2(x)\n",
    "        x = self.ff(x)\n",
    "        x = x + shortcut  # Add the original input back\n",
    "\n",
    "        return x, next_cache\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e88de3e3-9f07-42cc-816b-28dbd46e96c4",
   "metadata": {
    "id": "e88de3e3-9f07-42cc-816b-28dbd46e96c4"
   },
   "outputs": [],
   "source": [
    "class Qwen3Model(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "\n",
    "        # Main model parameters\n",
    "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"])\n",
    "\n",
    "        self.trf_blocks = nn.ModuleList(  # ModuleList since Sequential can only accept one input, and we need `x, mask, cos, sin`\n",
    "            [TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])]\n",
    "        )\n",
    "\n",
    "        self.final_norm = RMSNorm(cfg[\"emb_dim\"])\n",
    "        self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "\n",
    "        # Reusuable utilities\n",
    "        if cfg[\"head_dim\"] is None:\n",
    "            head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
    "        else:\n",
    "            head_dim = cfg[\"head_dim\"]\n",
    "        cos, sin = compute_rope_params(\n",
    "            head_dim=head_dim,\n",
    "            theta_base=cfg[\"rope_base\"],\n",
    "            context_length=cfg[\"context_length\"]\n",
    "        )\n",
    "        self.register_buffer(\"cos\", cos, persistent=False)\n",
    "        self.register_buffer(\"sin\", sin, persistent=False)\n",
    "        self.cfg = cfg\n",
    "        self.current_pos = 0  # Track current position in KV cache\n",
    "\n",
    "\n",
    "    def forward(self, in_idx, cache=None):\n",
    "        # Forward pass\n",
    "        tok_embeds = self.tok_emb(in_idx)\n",
    "        x = tok_embeds\n",
    "\n",
    "        num_tokens = x.shape[1]\n",
    "        if cache is not None:\n",
    "            pos_start = self.current_pos\n",
    "            pos_end = pos_start + num_tokens\n",
    "            self.current_pos = pos_end\n",
    "            mask = torch.triu(\n",
    "                torch.ones(pos_end, pos_end, device=x.device, dtype=torch.bool), diagonal=1\n",
    "            )[pos_start:pos_end, :pos_end]\n",
    "        else:\n",
    "            pos_start = 0  # Not strictly necessary but helps torch.compile\n",
    "            mask = torch.triu(\n",
    "                torch.ones(num_tokens, num_tokens, device=x.device, dtype=torch.bool), diagonal=1\n",
    "            )\n",
    "        # Shape (1, 1, num_tokens, num_tokens) to broadcast across batch and heads\n",
    "        mask = mask[None, None, :, :]\n",
    "\n",
    "        for i, block in enumerate(self.trf_blocks):\n",
    "            blk_cache = cache.get(i) if cache else None\n",
    "            x, new_blk_cache = block(x, mask, self.cos, self.sin,\n",
    "                                     start_pos=pos_start,\n",
    "                                     cache=blk_cache)\n",
    "            if cache is not None:\n",
    "                cache.update(i, new_blk_cache)\n",
    "\n",
    "        x = self.final_norm(x)\n",
    "        logits = self.out_head(x.to(self.cfg[\"dtype\"]))\n",
    "        return logits\n",
    "\n",
    "    def reset_kv_cache(self):\n",
    "        self.current_pos = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bc04d120",
   "metadata": {},
   "outputs": [],
   "source": [
    "class KVCache:\n",
    "    def __init__(self, n_layers):\n",
    "        self.cache = [None] * n_layers\n",
    "\n",
    "    def get(self, layer_idx):\n",
    "        return self.cache[layer_idx]\n",
    "\n",
    "    def update(self, layer_idx, value):\n",
    "        self.cache[layer_idx] = value\n",
    "\n",
    "    def get_all(self):\n",
    "        return self.cache\n",
    "\n",
    "    def reset(self):\n",
    "        for i in range(len(self.cache)):\n",
    "            self.cache[i] = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2d201f-74ad-4d63-ab9c-601b00674a48",
   "metadata": {
    "id": "be2d201f-74ad-4d63-ab9c-601b00674a48"
   },
   "source": [
    "&nbsp;\n",
    "# 2. Initialize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "caa142fa-b375-4e78-b392-2072ced666f3",
   "metadata": {
    "id": "caa142fa-b375-4e78-b392-2072ced666f3"
   },
   "outputs": [],
   "source": [
    "# Same config for\n",
    "\n",
    "# https://huggingface.co/Qwen/Qwen3-Coder-30B-A3B-Instruct (Qwen3 Coder Flash)\n",
    "# https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507\n",
    "# https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507\n",
    "# https://huggingface.co/Qwen/Qwen3-30B-A3B (original Instruct/Thinking hybrid model)\n",
    "\n",
    "QWEN3_CONFIG = {\n",
    "    \"vocab_size\": 151_936,\n",
    "    \"context_length\": 262_144,\n",
    "    \"emb_dim\": 2048,\n",
    "    \"n_heads\": 32,\n",
    "    \"n_layers\": 48,\n",
    "    \"head_dim\": 128,\n",
    "    \"qk_norm\": True,\n",
    "    \"n_kv_groups\": 4,\n",
    "    \"rope_base\": 10_000_000.0,\n",
    "    \"dtype\": torch.bfloat16,\n",
    "    \"num_experts\": 128,\n",
    "    \"num_experts_per_tok\": 8,\n",
    "    \"moe_hidden_dim\": 768,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "313effd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "156253fe-aacd-4da2-8f13-705f05c4b11e",
   "metadata": {
    "id": "156253fe-aacd-4da2-8f13-705f05c4b11e"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "with device:\n",
    "    model = Qwen3Model(QWEN3_CONFIG)\n",
    "\n",
    "#model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90aca91d-4bee-45ce-993a-4ec5393abe2b",
   "metadata": {},
   "source": [
    "- A quick check that the forward pass works before continuing (nan values are ok for now since we are using a \"meta\" device upon instantiation to save memory):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "adf0a6b7-b688-42c9-966e-c223d34db99f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[nan, nan, nan,  ..., nan, nan, nan],\n",
       "         [nan, nan, nan,  ..., nan, nan, nan],\n",
       "         [nan, nan, nan,  ..., nan, nan, nan]]], device='cuda:0',\n",
       "       dtype=torch.bfloat16, grad_fn=<UnsafeViewBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(torch.tensor([1, 2, 3]).unsqueeze(0).to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "364e76ca-52f8-4fa5-af37-c4069f9694bc",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "364e76ca-52f8-4fa5-af37-c4069f9694bc",
    "outputId": "00d7e983-262e-4c65-f322-f4d999311988"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of parameters: 30,532,122,624\n",
      "\n",
      "Total number of unique parameters: 30,220,957,696\n"
     ]
    }
   ],
   "source": [
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Total number of parameters: {total_params:,}\")\n",
    "\n",
    "# Account for weight tying\n",
    "total_params_normalized = total_params - model.tok_emb.weight.numel()\n",
    "print(f\"\\nTotal number of unique parameters: {total_params_normalized:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fd5efb03-5a07-46e8-8607-93ed47549d2b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fd5efb03-5a07-46e8-8607-93ed47549d2b",
    "outputId": "65c1a95e-b502-4150-9e2e-da619d9053d5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float32 (PyTorch default): 227.73 GB\n",
      "bfloat16: 113.87 GB\n"
     ]
    }
   ],
   "source": [
    "def model_memory_size(model, input_dtype=torch.float32):\n",
    "    total_params = 0\n",
    "    total_grads = 0\n",
    "    for param in model.parameters():\n",
    "        # Calculate total number of elements per parameter\n",
    "        param_size = param.numel()\n",
    "        total_params += param_size\n",
    "        # Check if gradients are stored for this parameter\n",
    "        if param.requires_grad:\n",
    "            total_grads += param_size\n",
    "\n",
    "    # Calculate buffer size (non-parameters that require memory)\n",
    "    total_buffers = sum(buf.numel() for buf in model.buffers())\n",
    "\n",
    "    # Size in bytes = (Number of elements) * (Size of each element in bytes)\n",
    "    # We assume parameters and gradients are stored in the same type as input dtype\n",
    "    element_size = torch.tensor(0, dtype=input_dtype).element_size()\n",
    "    total_memory_bytes = (total_params + total_grads + total_buffers) * element_size\n",
    "\n",
    "    # Convert bytes to gigabytes\n",
    "    total_memory_gb = total_memory_bytes / (1024**3)\n",
    "\n",
    "    return total_memory_gb\n",
    "\n",
    "print(f\"float32 (PyTorch default): {model_memory_size(model, input_dtype=torch.float32):.2f} GB\")\n",
    "print(f\"bfloat16: {model_memory_size(model, input_dtype=torch.bfloat16):.2f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4686eeb7-281f-4c5c-b37a-ed21d0a10427",
   "metadata": {},
   "source": [
    "- Don't be concerned; the model runs fine on an A100 card with 80 GB RAM due to offloading some layers to CPU RAM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c172f89f-d301-439f-b809-46169e5f5945",
   "metadata": {
    "id": "c172f89f-d301-439f-b809-46169e5f5945"
   },
   "source": [
    "&nbsp;\n",
    "# 4. Load pretrained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "75166128-5899-4995-9b88-9672e135650e",
   "metadata": {
    "id": "75166128-5899-4995-9b88-9672e135650e"
   },
   "outputs": [],
   "source": [
    "def load_weights_into_qwen(model, param_config, params):\n",
    "    def assign(left, right, tensor_name=\"unknown\"):\n",
    "        if left.shape != right.shape:\n",
    "            raise ValueError(f\"Shape mismatch in tensor '{tensor_name}'. Left: {left.shape}, Right: {right.shape}\")\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            if isinstance(right, torch.Tensor):\n",
    "                left.copy_(right)\n",
    "            else:\n",
    "                left.copy_(torch.as_tensor(right, dtype=left.dtype, device=left.device))\n",
    "    \n",
    "        return left \n",
    "\n",
    "    model.tok_emb.weight = assign(model.tok_emb.weight, params[\"model.embed_tokens.weight\"], \"model.embed_tokens.weight\")\n",
    "\n",
    "    for l in range(param_config[\"n_layers\"]):\n",
    "        block = model.trf_blocks[l]\n",
    "        att = block.att\n",
    "\n",
    "        # Q, K, V projections\n",
    "        att.W_query.weight = assign(\n",
    "            att.W_query.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.q_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.q_proj.weight\"\n",
    "        )\n",
    "        att.W_key.weight = assign(\n",
    "            att.W_key.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.k_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.k_proj.weight\"\n",
    "        )\n",
    "        att.W_value.weight = assign(\n",
    "            att.W_value.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.v_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.v_proj.weight\"\n",
    "        )\n",
    "\n",
    "        # Output projection\n",
    "        att.out_proj.weight = assign(\n",
    "            att.out_proj.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.o_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.o_proj.weight\"\n",
    "        )\n",
    "\n",
    "        # QK norms\n",
    "        if hasattr(att, \"q_norm\") and att.q_norm is not None:\n",
    "            att.q_norm.scale = assign(\n",
    "                att.q_norm.scale,\n",
    "                params[f\"model.layers.{l}.self_attn.q_norm.weight\"],\n",
    "                f\"model.layers.{l}.self_attn.q_norm.weight\"\n",
    "            )\n",
    "        if hasattr(att, \"k_norm\") and att.k_norm is not None:\n",
    "            att.k_norm.scale = assign(\n",
    "                att.k_norm.scale,\n",
    "                params[f\"model.layers.{l}.self_attn.k_norm.weight\"],\n",
    "                f\"model.layers.{l}.self_attn.k_norm.weight\"\n",
    "            )\n",
    "\n",
    "        # Attention layernorm\n",
    "        block.norm1.scale = assign(\n",
    "            block.norm1.scale,\n",
    "            params[f\"model.layers.{l}.input_layernorm.weight\"],\n",
    "            f\"model.layers.{l}.input_layernorm.weight\"\n",
    "        )\n",
    "\n",
    "        # Feedforward weights\n",
    "        if \"num_experts\" in param_config and param_config[\"num_experts\"] > 0:\n",
    "            # Load router (gating) weights\n",
    "            block.ff.gate.weight = assign(\n",
    "                block.ff.gate.weight,\n",
    "                params[f\"model.layers.{l}.mlp.gate.weight\"],\n",
    "                f\"model.layers.{l}.mlp.gate.weight\"\n",
    "            )\n",
    "            # Load expert weights\n",
    "            for e in range(param_config[\"num_experts\"]):\n",
    "                prefix = f\"model.layers.{l}.mlp.experts.{e}\"\n",
    "                block.ff.fc1[e].weight = assign(\n",
    "                    block.ff.fc1[e].weight,\n",
    "                    params[f\"{prefix}.gate_proj.weight\"],\n",
    "                    f\"{prefix}.gate_proj.weight\"\n",
    "                )\n",
    "                block.ff.fc2[e].weight = assign(\n",
    "                    block.ff.fc2[e].weight,\n",
    "                    params[f\"{prefix}.up_proj.weight\"],\n",
    "                    f\"{prefix}.up_proj.weight\"\n",
    "                )\n",
    "                block.ff.fc3[e].weight = assign(\n",
    "                    block.ff.fc3[e].weight,\n",
    "                    params[f\"{prefix}.down_proj.weight\"],\n",
    "                    f\"{prefix}.down_proj.weight\"\n",
    "                )\n",
    "\n",
    "        else:\n",
    "            block.ff.fc1.weight = assign(\n",
    "                block.ff.fc1.weight,\n",
    "                params[f\"model.layers.{l}.mlp.gate_proj.weight\"],\n",
    "                f\"model.layers.{l}.mlp.gate_proj.weight\"\n",
    "            )\n",
    "            block.ff.fc2.weight = assign(\n",
    "                block.ff.fc2.weight,\n",
    "                params[f\"model.layers.{l}.mlp.up_proj.weight\"],\n",
    "                f\"model.layers.{l}.mlp.up_proj.weight\"\n",
    "            )\n",
    "            block.ff.fc3.weight = assign(\n",
    "                block.ff.fc3.weight,\n",
    "                params[f\"model.layers.{l}.mlp.down_proj.weight\"],\n",
    "                f\"model.layers.{l}.mlp.down_proj.weight\"\n",
    "            )\n",
    "\n",
    "        block.norm2.scale = assign(\n",
    "            block.norm2.scale,\n",
    "            params[f\"model.layers.{l}.post_attention_layernorm.weight\"],\n",
    "            f\"model.layers.{l}.post_attention_layernorm.weight\"\n",
    "        )\n",
    "\n",
    "    # Final normalization and output head\n",
    "    model.final_norm.scale = assign(model.final_norm.scale, params[\"model.norm.weight\"], \"model.norm.weight\")\n",
    "\n",
    "    if \"lm_head.weight\" in params:\n",
    "        model.out_head.weight = assign(model.out_head.weight, params[\"lm_head.weight\"], \"lm_head.weight\")\n",
    "    else:\n",
    "        model.out_head.weight = model.tok_emb.weight\n",
    "        print(\"Model uses weight tying.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "699cb1b8-a67d-49fb-80a6-0dad9d81f392",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17,
     "referenced_widgets": [
      "9881b6995c3f49dc89e6992fd9ab660b",
      "17a3174e65c54476b2e0d1faf8f011ca",
      "1bbf2e62c0754d1593beb4105a7f1ac1",
      "b82112e1dec645d98aa1c1ba64abcb61",
      "271e2bd6a35e4a8b92de8697f7c0be5f",
      "90a79523187446dfa692723b2e5833a7",
      "431ffb83b8c14bf182f0430e07ea6154",
      "a8f1b72a33dd4b548de23fbd95e0da18",
      "25cc36132d384189acfbecc59483134b",
      "bfd06423ad544218968648016e731a46",
      "d029630b63ff44cf807ade428d2eb421"
     ]
    },
    "id": "699cb1b8-a67d-49fb-80a6-0dad9d81f392",
    "outputId": "55b2f28c-142f-4698-9d23-d27456d3ed6d"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "acdfb3a707444d7691bc8f1b053224b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 27 files:   0%|          | 0/27 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "from pathlib import Path\n",
    "from safetensors.torch import load_file\n",
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "repo_id = \"Qwen/Qwen3-30B-A3B\"  # Original Instruct/Thinking hybrind model\n",
    "repo_id = \"Qwen/Qwen3-235B-A22B-Instruct-2507\"  # New instruct model\n",
    "repo_id = \"Qwen/Qwen3-30B-A3B-Thinking-2507\"  # New thinking model\n",
    "repo_id = \"Qwen/Qwen3-Coder-30B-A3B-Instruct\"  # (Qwen3 Coder Flash)\n",
    "\n",
    "local_dir = Path(repo_id).parts[-1]\n",
    "\n",
    "repo_dir = snapshot_download(repo_id=repo_id, local_dir=local_dir)\n",
    "index_path = os.path.join(repo_dir, \"model.safetensors.index.json\")\n",
    "with open(index_path, \"r\") as f:\n",
    "    index = json.load(f)\n",
    "\n",
    "weights_dict = {}\n",
    "for filename in set(index[\"weight_map\"].values()):\n",
    "    shard_path = os.path.join(repo_dir, filename)\n",
    "    shard = load_file(shard_path)\n",
    "    weights_dict.update(shard)\n",
    "\n",
    "load_weights_into_qwen(model, QWEN3_CONFIG, weights_dict)\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b345491-3510-4397-92d3-cd0a3fa3deee",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "# 4. Load tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b68ab489-48e5-471e-a814-56cda2d60f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from tokenizers import Tokenizer\n",
    "\n",
    "class Qwen3Tokenizer:\n",
    "    _SPECIALS = [\n",
    "        \"<|endoftext|>\",\n",
    "        \"<|im_start|>\", \"<|im_end|>\",\n",
    "        \"<|object_ref_start|>\", \"<|object_ref_end|>\",\n",
    "        \"<|box_start|>\", \"<|box_end|>\",\n",
    "        \"<|quad_start|>\", \"<|quad_end|>\",\n",
    "        \"<|vision_start|>\", \"<|vision_end|>\",\n",
    "        \"<|vision_pad|>\", \"<|image_pad|>\", \"<|video_pad|>\",\n",
    "        \"<think>\", \"</think>\"\n",
    "    ]\n",
    "    _SPLIT_RE = re.compile(r\"(<\\|[^>]+?\\|>|<think>|</think>)\")\n",
    "\n",
    "    def __init__(self, tokenizer_file_path=\"tokenizer.json\", repo_id=None,\n",
    "                 apply_chat_template=True, add_generation_prompt=False, add_thinking=False):\n",
    "\n",
    "        self.apply_chat_template = apply_chat_template\n",
    "        self.add_generation_prompt = add_generation_prompt\n",
    "        self.add_thinking = add_thinking\n",
    "\n",
    "        tok_file = Path(tokenizer_file_path)\n",
    "        self._tok = Tokenizer.from_file(str(tok_file))\n",
    "        self._special_to_id = {}\n",
    "        for t in self._SPECIALS:\n",
    "            tid = self._tok.token_to_id(t)\n",
    "            if tid is not None:\n",
    "                self._special_to_id[t] = tid\n",
    "\n",
    "        self.pad_token_id = self._special_to_id[\"<|endoftext|>\"]\n",
    "        self.eos_token_id = self.pad_token_id\n",
    "\n",
    "        if repo_id and \"Base\" not in repo_id:\n",
    "            eos_token = \"<|im_end|>\"\n",
    "        else:\n",
    "            eos_token = \"<|endoftext|>\"\n",
    "        if eos_token in self._special_to_id:\n",
    "            self.eos_token_id = self._special_to_id[eos_token]\n",
    "\n",
    "    def encode(self, text, chat_wrapped=None):\n",
    "        if chat_wrapped is None:\n",
    "            chat_wrapped = self.apply_chat_template\n",
    "\n",
    "        stripped = text.strip()\n",
    "        if stripped in self._special_to_id and \"\\n\" not in stripped:\n",
    "            return [self._special_to_id[stripped]]\n",
    "\n",
    "        if chat_wrapped:\n",
    "            text = self._wrap_chat(text)\n",
    "\n",
    "        ids = []\n",
    "        for part in filter(None, self._SPLIT_RE.split(text)):\n",
    "            if part in self._special_to_id:\n",
    "                ids.append(self._special_to_id[part])\n",
    "            else:\n",
    "                ids.extend(self._tok.encode(part).ids)\n",
    "        return ids\n",
    "\n",
    "    def decode(self, ids):\n",
    "        return self._tok.decode(ids, skip_special_tokens=False)\n",
    "\n",
    "    def _wrap_chat(self, user_msg):\n",
    "        s = f\"<|im_start|>user\\n{user_msg}<|im_end|>\\n\"\n",
    "        if self.add_generation_prompt:\n",
    "            s += \"<|im_start|>assistant\"\n",
    "            if self.add_thinking:\n",
    "                s += \"\\n\"\n",
    "            else:\n",
    "                s += \"\\n<think>\\n\\n</think>\\n\\n\"\n",
    "        return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7b6df8bc-7308-468e-93ce-2d5529ea7866",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_file_path = f\"{Path(repo_id).parts[-1]}/tokenizer.json\"\n",
    "\n",
    "tokenizer = Qwen3Tokenizer(\n",
    "    tokenizer_file_path=tokenizer_file_path,\n",
    "    repo_id=repo_id,\n",
    "    apply_chat_template=True,\n",
    "    add_generation_prompt=True,\n",
    "    add_thinking=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1946b534-e3af-431a-a222-391a60bfa892",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|im_start|>user\\nImplement a binary search function in Python<|im_end|>\\n<|im_start|>assistant\\n'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# prompt = \"Give me a short introduction to large language models.\"\n",
    "prompt = \"Implement a binary search function in Python\"\n",
    "\n",
    "\n",
    "input_token_ids = tokenizer.encode(prompt)\n",
    "text = tokenizer.decode(input_token_ids)\n",
    "text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57d07df1-4401-4792-b549-7c4cc5632323",
   "metadata": {
    "id": "57d07df1-4401-4792-b549-7c4cc5632323"
   },
   "source": [
    "&nbsp;\n",
    "# 5. Generate text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "60b9fc72",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_text_basic_stream(model, token_ids, max_new_tokens, eos_token_id=None, context_size=None):\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        cache = KVCache(n_layers=model.cfg[\"n_layers\"])\n",
    "        model.reset_kv_cache()\n",
    "\n",
    "        # Prime the cache with the initial context\n",
    "        logits = model(token_ids, cache=cache)\n",
    "\n",
    "        for _ in range(max_new_tokens):\n",
    "            next_token = torch.argmax(logits[:, -1], dim=-1, keepdim=True)\n",
    "\n",
    "            if eos_token_id is not None and torch.all(next_token == eos_token_id):\n",
    "                break\n",
    "\n",
    "            yield next_token\n",
    "\n",
    "            token_ids = torch.cat([token_ids, next_token], dim=1)\n",
    "\n",
    "            # Feed only the new token to the model; cache handles history\n",
    "            logits = model(next_token, cache=cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a5b30753",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here's a comprehensive implementation of binary search in Python with both iterative and recursive approaches:\n",
      "\n",
      "## Iterative Binary Search\n",
      "\n",
      "```python\n",
      "def binary_search(arr, target):\n",
      "    \"\"\"\n",
      "    Iterative binary search implementation\n",
      "    \n",
      "    Args:\n",
      "        arr: Sorted list of elements\n",
      "        target: Element to search for\n",
      "    \n",
      "    Returns:\n",
      "        int: Index of target if found, -1 if not found\n",
      "    \n",
      "    Time Complexity: O(log n)\n",
      "    Space Complexity: O(1)\n",
      "    \"\"\"\n",
      "    left = 0\n",
      "    right = len(arr) - 1\n",
      "    \n",
      "    while left <= right:\n",
      "        # Calculate middle index (avoiding potential overflow)\n",
      "        mid = left + (right - left) // 2\n",
      "        \n",
      "        if arr[mid] == target:\n",
      "            return mid\n",
      "        elif arr[mid] < target:\n",
      "            left = mid + 1\n",
      "        else:\n",
      "            right = mid - 1\n",
      "    \n",
      "    return -1  # Target not found\n",
      "```\n",
      "\n",
      "## Recursive Binary Search\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_token_ids_tensor = torch.tensor(input_token_ids, device=device).unsqueeze(0)\n",
    "\n",
    "\n",
    "for token in generate_text_basic_stream(\n",
    "    model=model,\n",
    "    token_ids=input_token_ids_tensor,\n",
    "    max_new_tokens=200,\n",
    "    eos_token_id=tokenizer.eos_token_id\n",
    "):\n",
    "    token_id = token.squeeze(0).tolist()\n",
    "    print(\n",
    "        tokenizer.decode(token_id),\n",
    "        end=\"\",\n",
    "        flush=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "549324d6-5c71-4147-ae21-2e67675faa3d",
   "metadata": {
    "id": "549324d6-5c71-4147-ae21-2e67675faa3d"
   },
   "source": [
    "&nbsp;\n",
    "# What's next?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6edaaae-2de1-406c-8ffa-897cdfa3808c",
   "metadata": {
    "id": "e6edaaae-2de1-406c-8ffa-897cdfa3808c"
   },
   "source": [
    "- Check out the [README.md](./README.md), to use this model via the `llms_from_scratch` package\n",
    "- For those interested in a comprehensive guide on building a large language model from scratch and gaining a deeper understanding of its mechanics, you might like my [Build a Large Language Model (From Scratch)](http://mng.bz/orYv)\n",
    "\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
